/*
 * Copyright © 2020-present zmzhou-star. All Rights Reserved.
 */

package com.zjh.admin.websocket.handle;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;


/**
 * WebSocket处理器
 *
 * @author zmzhou
 * @version 1.0
 * @title WebShellWebSocketHandler
 * @date 2021/2/22 20:58
 */
@Slf4j
@Component
public class GlobalWebSocketHandler extends TextWebSocketHandler {

	public static final ConcurrentHashMap<String, Set<WebSocketSession>> socketMap_userName_session = new ConcurrentHashMap<>();

	public static final ConcurrentHashMap<String, WebSocketSession> sessionMap_id_session = new ConcurrentHashMap<>();

	@Override
	protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
		String payload = message.getPayload();
		// 解析前端发送的消息
		if ("{\"type\":\"ping\"}".equals(payload)) {
			// 如果是心跳检测消息，则回复一个 pong
			session.sendMessage(new TextMessage("pong"));
		} else {
			// 其他消息处理逻辑...
			System.out.println("Received message: " + payload);
		}
	}

	@Override
	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
		Map<String, Object> attributes = session.getAttributes();
		Object _username = attributes.get("username");
		String username = "";
		if (!Objects.isNull(_username)) {
			username = (String) _username;
			Set<WebSocketSession> sessions = socketMap_userName_session.get(username);
			if (sessions == null) {

				Set<WebSocketSession> userSession = new HashSet<>();
				userSession.add(session);
				socketMap_userName_session.put(username, userSession);
				sessionMap_id_session.put(session.getId(), session);

			} else {
				sessions.add(session);
			}

		}


		System.out.println("webSocket 新连接: " + session.getId() + attributes.get("ipAddress") + "  " + username);
	}

	@Override
	public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {


		removeSessionById(session.getId(), socketMap_userName_session);
		sessionMap_id_session.remove(session.getId());

		System.out.println("webSocket 断开: " + session.getId() + ", Status: " + status);
	}

	private void removeSessionById(String sessionId, Map<String, Set<WebSocketSession>> sessionMap) {
		Iterator<Map.Entry<String, Set<WebSocketSession>>> mapIterator = sessionMap.entrySet().iterator();

		while (mapIterator.hasNext()) {
			Map.Entry<String, Set<WebSocketSession>> entry = mapIterator.next();
			Set<WebSocketSession> sessions = entry.getValue();

			Iterator<WebSocketSession> sessionIterator = sessions.iterator();
			while (sessionIterator.hasNext()) {
				if (sessionIterator.next().getId().equals(sessionId)) {
					sessionIterator.remove();
					if (sessions.isEmpty()) {
						mapIterator.remove();
					}
					return; // 找到并移除后立即返回
				}
			}
		}
	}
}
